import logging

import VoltWatcherPack.VoltTarget as VT
import VoltWatcherPack.Pointer90 as P90
import onnxruntime
import cv2
import numpy as np
import json

# 你需要的类别名称
CLASSES = ['Meter']

# 创建一个日志器
logger = logging.getLogger(__name__)



# 载入两个功能的函数
# 假设 detect_objects 和 draw 是检测仪表函数
# 假设 process_meter_image 是识别仪表角度的函数
def detect_objects(img, model_path, conf_thres=0.25, iou_thres=0.45, image_width=1920, image_height=4096):
    # 使用onnxruntime加载ONNX模型
    session = onnxruntime.InferenceSession(model_path)

    # 将输入图像调整为模型期望的大小
    original_img = cv2.resize(img, (image_width, image_height))

    # 将图像从HWC格式转换为CHW格式 (3, 1024, 1280)，并进行BGR到RGB的通道反转
    input_img = original_img[:, :, ::-1].transpose(2, 0, 1)

    # 将像素值归一化到0-1之间，并转换为float32类型
    input_img = input_img.astype(np.float32) / 255.0

    # 扩展维度，符合模型输入的格式 (1, 3, 1024, 1280)
    input_img = np.expand_dims(input_img, axis=0)

    # 使用ONNX模型进行推理，获取检测结果
    outputs = session.run(None, {session.get_inputs()[0].name: input_img})

    # 预测结果
    predictions = outputs[0]

    # 使用VT.filter_box过滤掉低于置信度阈值且重叠度超过IoU阈值的框
    outbox = VT.filter_box(predictions, conf_thres, iou_thres)

    # 如果没有检测到目标物体
    if len(outbox) == 0:
        logging.info('未找到目标物体。')
        return None, original_img
    else:
        # 如果检测到目标物体，绘制检测框并返回
        logging.info('找到目标物体。')
        dest_img = VT.draw(original_img, outbox)
        return outbox, dest_img


# 假设 process_meter_image 是一个处理仪表的图像并识别指针的函数
def process_meter_image(cropped_meter_img):
    # 这里调用指针识别的代码, 返回角度和处理后的图像
    angle_deg, result_image = P90.process_meter_image(cropped_meter_img)
    return angle_deg, result_image


# 扩充图像
def resize_with_padding(image, target_width, target_height):
    # 获取原始图像尺寸
    h, w = image.shape[:2]

    # 计算需要填充的像素数
    top = 0  # 上边不填充
    bottom = target_height - h if target_height > h else 0  # 下边填充
    left = 0  # 左边不填充
    right = target_width - w if target_width > w else 0  # 右边填充

    # 执行填充操作，填充颜色为黑色
    resized_image = cv2.copyMakeBorder(
        image,
        top, bottom, left, right,
        borderType=cv2.BORDER_CONSTANT,
        value=(0, 0, 0)  # 黑色填充
    )

    return resized_image


# 去除无用像素
def crop_black_area(image):
    """裁剪图像中黑色区域"""
    # 转为灰度图，便于检测黑色区域
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 找到灰度图中非零的区域，代表有图像内容的区域
    _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
    # 找到非黑色区域的边界
    x, y, w, h = cv2.boundingRect(thresh)
    # 裁剪图像，去除黑色区域
    cropped_image = image[y:y + h, x:x + w]
    return cropped_image


def detect_and_process_meters(image, model_path, conf_thres=0.25, iou_thres=0.45, image_width=1920, image_height=4096):
    # 存储每个仪表的角度信息
    meter_readings = []

    # 步骤1: 检测仪表
    image = resize_with_padding(image, image_width ,image_height)

    outbox, detected_img = detect_objects(image, model_path, conf_thres, iou_thres, image_width, image_height)

    detected_img = crop_black_area(detected_img)

    if outbox is None:
        logging.info("未检测到仪表。")
        return meter_readings  # 返回空列表

    # 步骤2: 遍历每个检测到的仪表，进行裁剪并识别指针角度
    for box in outbox:
        x1, y1, x2, y2 = box[:4].astype(int)  # 获取边界框坐标
        meter_crop = detected_img[y1:y2, x1:x2]  # 裁剪图像

        # 检查裁剪图像的高度，若小于500则进行等比例放大
        height = meter_crop.shape[0]

        if height < 50:
            break

        logging.info(f"当前图像高度为: {height}")

        if height < 500:
            scale_factor = 500 / height
            new_width = int(meter_crop.shape[1] * scale_factor)
            meter_crop = cv2.resize(meter_crop, (new_width, 500))  # 等比例放大，保持高度为500

        # 调用指针识别函数
        angle_deg, result_image = process_meter_image(meter_crop)

        # 将每个仪表的角度和裁剪后的图像存储到列表中
        meter_readings.append({
            "angle": angle_deg,
            "meter_crop": meter_crop,
            "processed_image": result_image,
            "x1": x1,  # 添加x1坐标用于排序
            "y1": y1  # 添加y1坐标用于排序
        })

        # 显示结果
        if angle_deg is not None:
            logging.info(f"识别到仪表的角度: {angle_deg:.2f}度")
            # 在原始图像上绘制角度
            cv2.putText(detected_img, f"Angle: {angle_deg:.2f} deg", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    # 定义y轴方向的接近范围
    y_tolerance = 20  # 设置接近范围，例如20像素

    # 对y1进行归一化处理
    for reading in meter_readings:
        # 根据接近范围对y1值进行分组（向下取整到范围的倍数）
        reading['y1_grouped'] = reading['y1'] // y_tolerance

    # 按y1_grouped和x1排序
    meter_readings = sorted(
        meter_readings,
        key=lambda x: (x['y1_grouped'], x['x1'])  # 先按y1分组，再按x1排序
    )

    return meter_readings, detected_img  # 返回排序后的仪表角度信息


# 示例调用
if __name__ == "__main__":

    image_path = r"D:\Users\Administrator\Desktop\VoltWatcherDemo\stitched_result.jpg"
    model_path = r"D:\Users\Administrator\Desktop\VoltWatcherDemo\Meter1920x4096.onnx"

    image = cv2.imread(image_path)

    # 运行检测和识别函数
    meter_readings, detected_img = detect_and_process_meters(image, model_path, conf_thres=0.5, iou_thres=0.4)

    # 打印和显示每个仪表的角度和图像
    for i, reading in enumerate(meter_readings):
        logging.info(f"仪表 {i + 1} 的角度: {reading['angle']:.2f} 度")

        # 显示每个仪表的裁剪图像
        cv2.imshow(f"Meter {i + 1} Crop", reading["meter_crop"])

        # 显示每个仪表经过识别后的图像
        cv2.imshow(f"Meter {i + 1} Processed", reading["processed_image"])

        # 等待用户按键，之后再显示下一个仪表
        cv2.waitKey(0)

    # 最后显示带有所有检测框和角度的最终图像
    cv2.imshow("Final Image with Angles", detected_img)
    cv2.imwrite("detect_and_process_meters-out.jpg", detected_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def process_meter_readings(
        list_input,
        warehouse_key,
        image,
        model_path,
        conf_thres=0.5,
        iou_thres=0.4
):
    """
    处理指定仓库键的仪表读数，并更新数据列表。

    :param list_input: 包含多个字典的输入数据列表
    :param warehouse_key: 要筛选的仓库键
    :param image_path: 图像路径
    :param model_path: 模型路径
    :param conf_thres: 置信度阈值
    :param iou_thres: IOU阈值
    :return: 更新后的数据列表
    """

    # 定义过滤函数
    def filter_by_key(data_list, key, value):
        return [item for item in data_list if item[key] == value]

    # 定义提取位置信息函数
    def extract_positions(filtered_data):
        return [json.loads(item['position']) for item in filtered_data]

    # 定义处理位置信息函数
    def process_positions(positions):
        return sorted(set(int(position['Y']) for position in positions))

    # 将输入数据复制一份作为输出数据的基础
    list_output = list_input.copy()

    # 获取指定仓库键的数据
    result = filter_by_key(list_input, 'warehouse_key', warehouse_key)

    # 提取位置信息
    positions = extract_positions(result)

    # 处理位置信息，提取Y值并排序
    sorted_y_values = process_positions(positions)
    layer_num = len(sorted_y_values)

    # 遍历每个高度并处理图像
    for layer in range(layer_num):
        y_value = sorted_y_values[layer]
        logging.info(f"现在的高度是: {y_value}")

        # 调用VoltWatcher的检测处理函数
        meter_readings, detected_img = detect_and_process_meters(
            image, model_path, conf_thres=conf_thres, iou_thres=iou_thres
        )

        # 检查是否检测到了仪表
        if not meter_readings:
            logging.info(f"警告：在高度 {y_value} 没有检测到任何仪表")
            continue

        # 更新输出数据中的仪表读数
        for i in range(len(meter_readings)):
            logging.info(f"第 {i} 块表的值为 {meter_readings[i]['angle']:.2f}")

            # 遍历输出数据列表，匹配位置和表号
            for item in list_output:
                position_dict = json.loads(item['position'])
                item_y = int(position_dict['Y'])
                target_ptl_bm = f"b{layer}{i + 1}"

                # 基于Y值和ptl_bm匹配
                if item['ptl_bm'] == target_ptl_bm and item_y == y_value:
                    item['ptlValue'] = f"{meter_readings[i]['angle']:.2f}"
                    logging.info(f"更新了 {item['ptl_bm']} 的值为 {item['ptlValue']}")
                    break

    # 返回更新后的数据
    return list_output
